package com.binance.graceful.shutdown.core;

import org.eclipse.jetty.server.Server;
import org.eclipse.jetty.server.handler.StatisticsHandler;
import org.springframework.boot.context.embedded.jetty.JettyServerCustomizer;
import org.springframework.context.ApplicationListener;
import org.springframework.context.event.ContextClosedEvent;

import lombok.extern.log4j.Log4j2;

/**
 * Created by robin.wu on 2018/10/26.
 */
@Log4j2
public class GracefulShutdownJetty implements JettyServerCustomizer, ApplicationListener<ContextClosedEvent> {

    private volatile Server server;

    private Long waitTime;

    public GracefulShutdownJetty(Long waitTime){
        this.waitTime = waitTime;
    }

    @Override
    public void customize(Server server) {
        this.server = server;
        log.info("Jetty Server customizer,wait time:{}",this.waitTime);
        if (this.waitTime > 0) {
            StatisticsHandler handler = new StatisticsHandler();
            handler.setHandler(server.getHandler());
            server.setHandler(handler);

            server.setStopTimeout(this.waitTime);
            // We will stop it through JettyGracefulShutdown class.
            server.setStopAtShutdown(false);
        }
    }

    @Override
    public void onApplicationEvent(ContextClosedEvent event) {
        if (server == null) {
            log.error("Jetty server variable is null, this should not happen!");
            return;
        }
        log.info("Entering shutdown for Jetty.");
        if (!(server.getHandler() instanceof StatisticsHandler)) {
            log.error("Root handler is not StatisticsHandler, graceful shutdown may not work at all!");
        } else {
            log.info("Active requests: " + ((StatisticsHandler) server.getHandler()).getRequestsActive());
        }
        try {
            long begin = System.currentTimeMillis();
            server.stop();
            log.info("Shutdown took {} ms.",System.currentTimeMillis() - begin);
        } catch (Exception e) {
            log.error("Fail to shutdown gracefully.", e);
        }
    }


}
